// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/testing/gtest_utils.h>
#include <pollux/testing/exec/util/assert_query_builder.h>
#include <pollux/functions/lib/aggregates/tests/utils/AggregationTestBase.h>
#include <pollux/functions/prestosql/types/timestamp_with_time_zone_type.h>

using namespace kumo::pollux::exec;
using namespace kumo::pollux::plan;
using namespace kumo::pollux::functions::aggregate::test;

namespace kumo::pollux::aggregate::test {
    namespace {
        constexpr int64_t kLongMax = std::numeric_limits<int64_t>::max();
        constexpr int64_t kLongMin = std::numeric_limits<int64_t>::min();
        constexpr int128_t kHugeMax = std::numeric_limits<int128_t>::max();
        constexpr int128_t kHugeMin = std::numeric_limits<int128_t>::min();

        class SetUnionTest : public AggregationTestBase {
        protected:
            void SetUp() override {
                AggregationTestBase::SetUp();
            }
        };

        TEST_F(SetUnionTest, global) {
            auto data = make_row_vector({
                make_array_vector<int32_t>({
                    {},
                    {1, 2, 3},
                    {1, 2},
                    {2, 3, 4, 5},
                    {6, 7},
                }),
            });

            auto expected = make_row_vector({
                make_array_vector<int32_t>({
                    {1, 2, 3, 4, 5, 6, 7},
                }),
            });

            testAggregations(
                {data}, {}, {"set_union(c0)"}, {"array_sort(a0)"}, {expected});

            // Null inputs.
            data = make_row_vector({
                make_array_vector_from_json<int32_t>({
                    "[]",
                    "[1, 2, null, 3]",
                    "[1, 2]",
                    "null",
                    "[2, 3, 4, null, 5]",
                    "[6, 7]",
                }),
            });

            expected = make_row_vector({
                make_array_vector_from_json<int32_t>({"[1, 2, 3, 4, 5, 6, 7, null]"}),
            });

            testAggregations(
                {data}, {}, {"set_union(c0)"}, {"array_sort(a0)"}, {expected});

            // All nulls arrays.
            data = make_row_vector({
                make_all_null_array_vector(10, INTEGER()),
            });

            expected = make_row_vector({
                // Empty array: [].
                make_array_vector<int32_t>({{}}),
            });

            testAggregations(
                {data}, {}, {"set_union(c0)"}, {"array_sort(a0)"}, {expected});

            // All nulls elements.
            data = make_row_vector({
                make_array_vector_from_json<int32_t>({
                    "[]",
                    "[null, null, null, null]",
                    "[null, null, null]",
                }),
            });

            expected = make_row_vector({
                make_array_vector_from_json<int32_t>({"[null]"}),
            });

            testAggregations(
                {data}, {}, {"set_union(c0)"}, {"array_sort(a0)"}, {expected});
        }

        TEST_F(SetUnionTest, groupBy) {
            auto data = make_row_vector({
                make_flat_vector<int16_t>({1, 1, 2, 2, 2, 1, 2, 1, 2, 1}),
                make_array_vector<int32_t>({
                    {},
                    {1, 2, 3}, // masked out
                    {10, 20},
                    {20, 30, 40},
                    {10, 50}, // masked out
                    {4, 2, 1, 5}, // masked out
                    {60, 20},
                    {5, 6},
                    {}, // masked out
                    {},
                }),
                make_flat_vector<bool>(
                    {true, false, true, true, false, false, true, true, false, true}),
            });

            auto expected = make_row_vector({
                make_flat_vector<int16_t>({1, 2}),
                make_array_vector<int32_t>({
                    {1, 2, 3, 4, 5, 6},
                    {10, 20, 30, 40, 50, 60},
                }),
            });

            testAggregations(
                {data, data, data},
                {"c0"},
                {"set_union(c1)"},
                {"c0", "array_sort(a0)"},
                {expected});

            expected = make_row_vector({
                make_flat_vector<int16_t>({1, 2}),
                make_array_vector<int32_t>({
                    {5, 6},
                    {10, 20, 30, 40, 60},
                }),
            });

            testAggregations(
                {data, data, data},
                {"c0"},
                {"set_union(c1) filter (where c2)"},
                {"c0", "array_sort(a0)"},
                {expected});

            // Null inputs.
            data = make_row_vector({
                make_flat_vector<int16_t>({1, 1, 2, 2, 2, 1, 2, 1, 2, 1}),
                make_array_vector_from_json<int32_t>({
                    "[]",
                    "[1, 2, 3]",
                    "[10, null, 20]",
                    "[20, 30, 40, null, 50]",
                    "null",
                    "[4, 2, 1, 5]",
                    "[60, null, 20]",
                    "[null, 5, 6]",
                    "[]",
                    "null",
                }),
            });

            expected = make_row_vector({
                make_flat_vector<int16_t>({1, 2}),
                make_array_vector_from_json<int32_t>({
                    "[1, 2, 3, 4, 5, 6, null]",
                    "[10, 20, 30, 40, 50, 60, null]",
                }),
            });

            testAggregations(
                {data, data, data},
                {"c0"},
                {"set_union(c1)"},
                {"c0", "array_sort(a0)"},
                {expected});

            // All null arrays for one group.
            std::vector<RowVectorPtr> multiBatchData = {
                make_row_vector({
                    make_flat_vector<int16_t>({1, 1, 2, 2, 2, 1, 2}),
                    make_array_vector_from_json<int32_t>({
                        "null",
                        "null",
                        "[]",
                        "[1, 2]",
                        "[1, 2, 3]",
                        "null",
                        "null",
                    }),
                }),
                make_row_vector({
                    make_flat_vector<int16_t>({3, 3, 3, 2, 3}),
                    make_array_vector_from_json<int32_t>({
                        "null",
                        "null",
                        "null",
                        "[2, 4, 5]",
                        "null",
                    }),
                }),
            };

            expected = make_row_vector({
                make_flat_vector<int16_t>({1, 2, 3}),
                make_array_vector<int32_t>({
                    {}, // Empty array: [].
                    {1, 2, 3, 4, 5},
                    {}, // Empty array: [].
                }),
            });

            testAggregations(
                multiBatchData,
                {"c0"},
                {"set_union(c1)"},
                {"c0", "array_sort(a0)"},
                {expected});
        }

        TEST_F(SetUnionTest, inputOrder) {
            // Presto preserves order of input.

            auto testInputOrder = [&](const RowVectorPtr &data,
                                      const RowVectorPtr &expected) {
                auto plan = PlanBuilder()
                        .values({data})
                        .singleAggregation({}, {"set_union(c0)"})
                        .planNode();
                assertQuery(plan, expected);
            };

            auto data = make_row_vector({
                make_array_vector_from_json<int32_t>({
                    "[]",
                    "[1, 2, 3]",
                    "[1, 2, null]",
                    "[2, 3, 4, 5]",
                    "[6, 7]",
                }),
            });

            auto expected = make_row_vector({
                make_nullable_array_vector<int32_t>({
                    {1, 2, 3, std::nullopt, 4, 5, 6, 7},
                }),
            });

            testInputOrder(data, expected);

            // Strings.
            data = make_row_vector({
                make_nullable_array_vector<StringView>({
                    {},
                    {"abc", "bxy", "cde"},
                    {"abc", "bxy"},
                    {"cdef", "hijk", std::nullopt},
                    {"abc", "some very long string to test long strings"},
                }),
            });

            expected = make_row_vector({
                make_nullable_array_vector<StringView>({
                    {
                        "abc",
                        "bxy",
                        "cde",
                        "cdef",
                        "hijk",
                        std::nullopt,
                        "some very long string to test long strings"
                    },
                }),
            });

            testInputOrder(data, expected);

            // Complex types.

            data = make_row_vector({
                make_nested_array_vector_from_json<int32_t>(
                    {"[[1,2], [5, 6], null]", "[[3,4], [7, 8], null]"})
            });

            expected = make_row_vector({
                make_nested_array_vector_from_json<int32_t>(
                    {"[[1,2], [5, 6], null, [3,4], [7, 8]]"})
            });

            testInputOrder(data, expected);

            // Group by.
            data = make_row_vector({
                make_flat_vector<int16_t>({
                    1,
                    2,
                    1,
                    2,
                    1,
                }),
                make_array_vector_from_json<int32_t>({
                    "[]",
                    "[1, 2, 3]",
                    "[1, 2, null]",
                    "[2, 3, 4, 5]",
                    "[6, 7]",
                }),
            });

            expected = make_row_vector({
                make_flat_vector<int16_t>({1, 2}),
                make_array_vector_from_json<int32_t>({
                    "[1, 2, null, 6, 7]",
                    "[1, 2, 3, 4, 5]",
                }),
            });

            auto plan = PlanBuilder()
                    .values({data})
                    .singleAggregation({"c0"}, {"set_union(c1)"})
                    .planNode();

            assertQuery(plan, expected);
        }

        TEST_F(SetUnionTest, shortDecimal) {
            // Test with short decimal
            auto type = DECIMAL(6, 2);

            auto data = make_row_vector({
                make_array_vector<int64_t>(
                    {
                        {},
                        {kLongMin, 2000, 3000},
                        {kLongMin, -2000},
                        {2000, 3000, kLongMax, 5000},
                        {6000, 7000, -5432},
                    },
                    type),
            });

            auto expected = make_row_vector({
                make_array_vector<int64_t>(
                    {
                        {kLongMin, -5432, -2000, 2000, 3000, 5000, 6000, 7000, kLongMax},
                    },
                    type),
            });

            testAggregations(
                {data}, {}, {"set_union(c0)"}, {"array_sort(a0)"}, {expected});

            // Test with some NULL inputs (short decimals)
            data = make_row_vector({
                make_nullable_array_vector<int64_t>(
                    {
                        {},
                        {-1000, std::nullopt, kLongMax, std::nullopt, 7000},
                        {1000, kLongMax, std::nullopt, 4000, std::nullopt, std::nullopt},
                        {kLongMin, -5923},
                    },
                    ARRAY(type)),
            });

            expected = make_row_vector({
                make_nullable_array_vector(
                    std::vector<std::vector<std::optional<int64_t> > >{
                        {
                            kLongMin,
                            -5923,
                            -1000,
                            1000,
                            4000,
                            7000,
                            kLongMax,
                            std::nullopt
                        },
                    },
                    ARRAY(type)),
            });

            testAggregations(
                {data}, {}, {"set_union(c0)"}, {"array_sort(a0)"}, {expected});

            // Test with all NULL inputs (short decimals)
            data = make_row_vector({
                make_all_null_array_vector(10, type),
            });

            expected = make_row_vector({
                make_array_vector<int64_t>({{}}, type),
            });

            testAggregations(
                {data}, {}, {"set_union(c0)"}, {"array_sort(a0)"}, {expected});

            data = make_row_vector({
                make_nullable_array_vector<int64_t>(
                    {
                        {},
                        {std::nullopt, std::nullopt, std::nullopt, std::nullopt},
                        {std::nullopt, std::nullopt, std::nullopt},
                    },
                    ARRAY(type)),
            });

            expected = make_row_vector({
                make_nullable_array_vector(
                    std::vector<std::vector<std::optional<int64_t> > >{
                        {std::nullopt},
                    },
                    ARRAY(type)),
            });

            testAggregations(
                {data}, {}, {"set_union(c0)"}, {"array_sort(a0)"}, {expected});
        }

        TEST_F(SetUnionTest, longDecimal) {
            // Test with long decimals
            auto type = DECIMAL(20, 2);

            auto data = make_row_vector({
                make_array_vector<int128_t>(
                    {
                        {},
                        {kHugeMax, -2000, 3000},
                        {1000, 2000},
                        {2000, kHugeMin, 3000, 4000, 5000},
                        {-6363, 7000},
                    },
                    type),
            });

            auto expected = make_row_vector({
                make_array_vector<int128_t>(
                    {
                        {
                            kHugeMin,
                            -6363,
                            -2000,
                            1000,
                            2000,
                            3000,
                            4000,
                            5000,
                            7000,
                            kHugeMax
                        },
                    },
                    type),
            });

            testAggregations(
                {data}, {}, {"set_union(c0)"}, {"array_sort(a0)"}, {expected});

            // Test with some NULL inputs (long decimals).
            data = make_row_vector({
                make_nullable_array_vector<int128_t>(
                    {
                        {},
                        {kHugeMin},
                        {1000, std::nullopt, 3000, std::nullopt, 7000},
                        {-1000, 3000, std::nullopt, 4000, std::nullopt, std::nullopt},
                        {2000, kHugeMax, -1234},
                    },
                    ARRAY(type)),
            });

            expected = make_row_vector({
                make_nullable_array_vector(
                    std::vector<std::vector<std::optional<int128_t> > >{
                        {
                            kHugeMin,
                            -1234,
                            -1000,
                            1000,
                            2000,
                            3000,
                            4000,
                            7000,
                            kHugeMax,
                            std::nullopt
                        },
                    },
                    ARRAY(type)),
            });

            testAggregations(
                {data}, {}, {"set_union(c0)"}, {"array_sort(a0)"}, {expected});

            // Test with all NULL inputs (long decimals).
            data = make_row_vector({
                make_all_null_array_vector(10, type),
            });

            expected = make_row_vector({
                make_array_vector<int128_t>({{}}, type),
            });

            testAggregations(
                {data}, {}, {"set_union(c0)"}, {"array_sort(a0)"}, {expected});
        }

        TEST_F(SetUnionTest, nans) {
            // Verify that NaNs with different binary representations are considered equal
            // and deduplicated.
            static const auto kNaN = std::numeric_limits<double>::quiet_NaN();
            static const auto kSNaN = std::numeric_limits<double>::signaling_NaN();

            // Global aggregation, Primitive type.
            auto data = make_row_vector({
                make_array_vector<double>({
                    {1, 2, kNaN},
                    {3, 4, kNaN},
                    {3, 2, kSNaN},
                }),
                make_flat_vector<int32_t>({1, 2, 2}),
            });

            auto expected = make_row_vector({
                make_array_vector<double>({
                    {1, 2, 3, 4, kNaN},
                }),
            });

            testAggregations(
                {data}, {}, {"set_union(c0)"}, {"array_sort(a0)"}, {expected});

            // Check again with kSNaN as the first NaN element to be encountered.
            data = make_row_vector({
                make_array_vector<double>({
                    {1, 2, kSNaN},
                    {3, 4, kSNaN},
                    {3, 2, kNaN},
                }),
                make_flat_vector<int32_t>({1, 2, 2}),
            });
            testAggregations(
                {data}, {}, {"set_union(c0)"}, {"array_sort(a0)"}, {expected});

            // Group by aggregation, Primitive type.
            expected = make_row_vector({
                make_array_vector<double>({
                    {1, 2, kNaN},
                    {2, 3, 4, kNaN},
                }),
                make_flat_vector<int32_t>({1, 2}),
            });

            testAggregations(
                {data}, {"c1"}, {"set_union(c0)"}, {"array_sort(a0)", "c1"}, {expected});

            // Global aggregation, Complex type.
            data = make_row_vector({
                make_array_vector(
                    {0, 3, 6},
                    make_row_vector({
                        make_flat_vector<StringView>({
                            "a"_sv,
                            "b"_sv,
                            "c"_sv,
                            "a"_sv,
                            "b"_sv,
                            "c"_sv,
                            "a"_sv,
                            "b"_sv,
                            "c"_sv,
                        }),
                        make_flat_vector<double>({1, 2, kNaN, 1, 2, kSNaN, 1, 2, kNaN}),
                    })),
                make_flat_vector<int32_t>({1, 2, 2}),
            });

            expected = make_row_vector({
                make_array_vector(
                    {0},
                    make_row_vector(
                        {
                            make_flat_vector<StringView>({
                                "a"_sv,
                                "b"_sv,
                                "c"_sv,
                            }),
                            make_flat_vector<double>({1, 2, kNaN})
                        }))
            });

            testAggregations(
                {data}, {}, {"set_union(c0)"}, {"array_sort(a0)"}, {expected});

            // Group by aggregation, Complex type.
            expected = make_row_vector({
                make_array_vector(
                    {0, 3},
                    make_row_vector(
                        {
                            make_flat_vector<StringView>({
                                "a"_sv,
                                "b"_sv,
                                "c"_sv,
                                "a"_sv,
                                "b"_sv,
                                "c"_sv,
                            }),
                            make_flat_vector<double>({1, 2, kNaN, 1, 2, kNaN})
                        })),
                make_flat_vector<int32_t>({1, 2}),
            });

            testAggregations(
                {data}, {"c1"}, {"set_union(c0)"}, {"array_sort(a0)", "c1"}, {expected});
        }

        TEST_F(SetUnionTest, TimestampWithTimeZone) {
            // Global aggregation, Primitive type.
            auto data = make_row_vector({
                make_array_vector<int64_t>(
                    {
                        {pack(0, 0), pack(1, 1), pack(2, 4)},
                        {pack(3, 0), pack(4, 0), pack(2, 8)},
                        {pack(3, 1), pack(1, 0), pack(2, 2)},
                    },
                    TIMESTAMP_WITH_TIME_ZONE()),
                make_flat_vector<int32_t>({1, 2, 2}),
            });

            auto expected = make_row_vector({
                make_array_vector<int64_t>(
                    {{pack(0, 0), pack(1, 1), pack(2, 4), pack(3, 0), pack(4, 0)}},
                    TIMESTAMP_WITH_TIME_ZONE()),
            });

            testAggregations(
                {data}, {}, {"set_union(c0)"}, {"array_sort(a0)"}, {expected});

            // Group by aggregation, Primitive type.
            expected = make_row_vector({
                make_array_vector<int64_t>(
                    {
                        {pack(0, 0), pack(1, 1), pack(2, 4)},
                        {pack(1, 0), pack(2, 8), pack(3, 0), pack(4, 0)},
                    },
                    TIMESTAMP_WITH_TIME_ZONE()),
                make_flat_vector<int32_t>({1, 2}),
            });

            testAggregations(
                {data}, {"c1"}, {"set_union(c0)"}, {"array_sort(a0)", "c1"}, {expected});

            // Global aggregation, wrapped in complex type.
            data = make_row_vector({
                make_array_vector(
                    {0, 3, 6},
                    make_row_vector({
                        make_flat_vector<int64_t>(
                            {
                                pack(0, 0),
                                pack(1, 0),
                                pack(2, 0),
                                pack(0, 1),
                                pack(1, 1),
                                pack(2, 1),
                                pack(0, 2),
                                pack(1, 2),
                                pack(2, 2)
                            },
                            TIMESTAMP_WITH_TIME_ZONE())
                    })),
                make_flat_vector<int32_t>({1, 2, 2}),
            });

            expected = make_row_vector({
                make_array_vector(
                    {0},
                    make_row_vector({
                        make_flat_vector<int64_t>(
                            {pack(0, 0), pack(1, 0), pack(2, 0)},
                            TIMESTAMP_WITH_TIME_ZONE())
                    }))
            });

            testAggregations(
                {data}, {}, {"set_union(c0)"}, {"array_sort(a0)"}, {expected});

            // Group by aggregation, wrapped in complex type.
            expected = make_row_vector({
                make_array_vector(
                    {0, 3},
                    make_row_vector({
                        make_flat_vector<int64_t>(
                            {
                                pack(0, 0),
                                pack(1, 0),
                                pack(2, 0),
                                pack(0, 1),
                                pack(1, 1),
                                pack(2, 1)
                            },
                            TIMESTAMP_WITH_TIME_ZONE())
                    })),
                make_flat_vector<int32_t>({1, 2}),
            });

            testAggregations(
                {data}, {"c1"}, {"set_union(c0)"}, {"array_sort(a0)", "c1"}, {expected});
        }

        TEST_F(SetUnionTest, unknownType) {
            // Test with unknown type
            auto type = UNKNOWN();

            auto data = make_row_vector({
                make_all_null_array_vector(10, type),
            });

            auto expected = make_row_vector({
                make_array_vector<int64_t>({{}}, type),
            });

            testAggregations(
                {data}, {}, {"set_union(c0)"}, {"array_sort(a0)"}, {expected});

            data = make_row_vector({
                make_nullable_array_vector<int64_t>(
                    {
                        {},
                        {std::nullopt, std::nullopt, std::nullopt, std::nullopt},
                        {std::nullopt, std::nullopt, std::nullopt},
                    },
                    ARRAY(type)),
            });

            expected = make_row_vector({
                make_nullable_array_vector(
                    std::vector<std::vector<std::optional<int64_t> > >{
                        {std::nullopt},
                    },
                    ARRAY(type)),
            });

            testAggregations(
                {data}, {}, {"set_union(c0)"}, {"array_sort(a0)"}, {expected});
        }
    } // namespace
} // namespace kumo::pollux::aggregate::test
